package com.byron.media.server.ffmpeg;

import com.byron.media.server.utils.FrameInfo;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.BytePointer;
import org.bytedeco.javacpp.DoublePointer;
import org.bytedeco.javacpp.avformat;

import static org.bytedeco.javacpp.avcodec.*;
import static org.bytedeco.javacpp.avutil.*;
import static org.bytedeco.javacpp.swscale.*;

/**
 * 任何视频源转换为固定高宽比例的H264裸流
 */
@Slf4j
public class AnyToH264 {

    private int width;
    private int height;
    private int dpi = 1;

    private AVCodecContext inCodecContext = null;           // 解码工具类
    private AVFrame inFrame = null;					// 刚解码的frame数据，未转换
    private AVPacket inPacket;
    private AVFrame outFrame = null;

    private FfmpegUtils ffmpegUtils;

    private SwsContext swsContext;

    private int videoDen;

    private int videoNum;

    private byte[] y;
    private byte[] u;
    private byte[] v;
    private FrameListener frameListener;

    public int handle(byte[] mediaData, long pts, long dts, long duration){
        int ret;
        inPacket.data().put(mediaData);
        inPacket.size(mediaData.length);
        ret = avcodec_send_packet(inCodecContext, inPacket);
        if (ret < 0) {
            inPacket.data().position(0);
            inPacket.size(0);
            log.info(String.format("Decode Error. %d ", ret));
            return -1;
        }
        ret = avcodec_receive_frame(inCodecContext, inFrame);
        if(ret == -11){         // AVERROR（EAGAIN）
            return 0;
        }

        if (ret < 0) {
            return -1;
        }

//        long inPts = inPacket.pts();
//        long inDts = inPacket.dts();
//        long inFamePts = inFrame.pts();
//        long inFameDts = inFrame.pkt_dts();


        // 解码出YUV
        int linesize0 = inFrame.linesize(0);
        int linesize1 = inFrame.linesize(1);
        int linesize2 = inFrame.linesize(2);
        int inWidth = inFrame.width();
        int inHeight = inFrame.height();
        int inFormat = inFrame.format();



        if(inFormat == AV_PIX_FMT_YUV420P){

            if(y == null){y = new byte[inWidth * inHeight]; }
            if(u == null){u = new byte[inWidth * inHeight / 4]; }
            if(v == null){v = new byte[inWidth * inHeight / 4]; }

            BytePointer yPointer = inFrame.data(0);
            yPointer.position(0).get(y);
            inFrame.data(1).position(0).get(u);
            inFrame.data(2).position(0).get(v);
//            inFrame.data(1).position(u.length).get(v);

            ffmpegUtils.pushYUV(y, u, v, AV_PIX_FMT_YUV420P, inWidth, inHeight, pts, dts, duration);
        } else {
            if(outFrame == null){
                outFrame = av_frame_alloc();
                outFrame
                        .format(AV_PIX_FMT_YUV420P)
                        .width(inWidth)
                        .height(inHeight);
                av_image_alloc(
                        outFrame.data(),
                        outFrame.linesize(),
                        outFrame.width(),
                        outFrame.height(),
                        AV_PIX_FMT_YUV420P,
                        16);
            }
            // 重新编码成YUV420P
            if(swsContext == null){
                swsContext = sws_getContext(
                        inWidth, inHeight, inFrame.format(),
                        inWidth, inHeight, outFrame.format(),
                        SWS_FAST_BILINEAR ,
//                        SWS_BICUBIC,
//                        SWS_SINC,
                        null, null, (DoublePointer) null);
                if (swsContext == null) {
                    log.info("Cannot initialize the conversion context");
                    return -1;
                }
            }

            ret = sws_scale(swsContext, inFrame.data(), inFrame.linesize(), 0, inHeight, outFrame.data(), outFrame.linesize());
            if(ret < 0){
                return -1;
            }
            if(y == null){y = new byte[inWidth * inHeight ]; }
            if(u == null){u = new byte[inWidth * inHeight / 4]; }
            if(v == null){v = new byte[inWidth * inHeight / 4]; }

            BytePointer yPointer = outFrame.data(0);
            yPointer.position(0).get(y);
            outFrame.data(1).position(0).get(u);
            outFrame.data(2).position(0).get(v);
            outFrame.pts(inFrame.pts());
            outFrame.pkt_dts(inFrame.pkt_dts());

            ffmpegUtils.pushYUV(y, u, v, AV_PIX_FMT_YUV420P, inWidth, inHeight, pts, dts, duration);
        }

        return 0;
    }

    /**
     * 初始化H264
     */
    public int initFfmpeg(avformat.AVFormatContext inFormatContext, int width, int height){
        this.width = width;
        this.height = height;
        int video_stream_index = -1;
        int audio_stream_index = -1;
        for (int i = 0; i < inFormatContext.nb_streams(); ++i) {
            avformat.AVStream stream = inFormatContext.streams(i);

            log.info(String.format("type of the encoded data: %d", stream.codecpar().codec_id()));
            if (stream.codecpar().codec_type() == AVMEDIA_TYPE_VIDEO) {
                video_stream_index = i;
                log.info(String.format("dimensions of the video frame in " +
                                "pixels: width: %d, height: %d, " +
                                "pixel format: %d, " +
                                "start time: %d," +
                                "time base: %d/%d",
                        stream.codecpar().width(), stream.codecpar().height(),
                        stream.codecpar().format(),
                        stream.start_time(),
                        stream.time_base().num(), stream.time_base().den()));
                this.videoDen = stream.time_base().den();
                this.videoNum = stream.time_base().num();

            } else if (stream.codecpar().codec_type() == AVMEDIA_TYPE_AUDIO) {
                audio_stream_index = i;
                log.info(String.format("audio sample format: %d", stream.codecpar().format()));
            }
        }

        if(video_stream_index == -1){
            return -1;
        }

        avformat.AVStream inStream = inFormatContext.streams(video_stream_index);

//        int inCodecId = AV_CODEC_ID_H264;
        int inCodecId = inStream.codecpar().codec_id();
        // 初始化解码器
        AVCodec inCodec = avcodec_find_decoder(inCodecId);
        if (inCodec == null) {
            log.info("Codec not found");
            return -1;
        }

        // 初始化解码器的上下文
//        inCodecContext = avcodec_alloc_context3(inCodec);
        inCodecContext = inStream.codec();
        if (inCodecContext == null){
            log.info("Could not allocate video codec context");
            return -1;
        }

        // 打开解码器
        if (avcodec_open2(inCodecContext, inCodec, (AVDictionary) null) < 0) {
            log.info("Could not open codec");
            return -1;
        }

        inFrame = av_frame_alloc();
        inPacket = av_packet_alloc();
        av_new_packet(inPacket, 1024 * 1024);


        // 初始化编码器
//        int outCodecId = AV_CODEC_ID_H264;
        AVCodec outCodec = avcodec_find_decoder(AV_CODEC_ID_H264);
        if (outCodec == null) {
            log.info("Codec not found");
            return -1;
        }

        ffmpegUtils = new FfmpegUtils();
        ffmpegUtils.encodeInit(width, height, 1, 10);
        ffmpegUtils.setFrameListener(frameInfo -> {
            if(AnyToH264.this.frameListener != null){
                AnyToH264.this.frameListener.onFrame(frameInfo);
            }
        });

        return 0;
    }


    /**
     * 释放资源
     * @return
     */
    public int releaseFfmpeg(){

        ffmpegUtils.encoderRelease();

        if(inCodecContext != null){
            avcodec_free_context(inCodecContext);
            inCodecContext.setNull();
            inCodecContext = null;
        }

        if(inFrame != null){
            av_frame_free(inFrame);
            inFrame.setNull();
            inFrame = null;
        }

        if(inPacket != null){
            av_packet_free(inPacket);
            inPacket.setNull();
            inPacket = null;
        }

        if(outFrame != null){
            av_frame_free(outFrame);
            outFrame.setNull();
            outFrame = null;
        }

        if (swsContext != null) {
            sws_freeContext(swsContext);
            swsContext.setNull();
            swsContext = null;
        }

        return 0;
    }
//    @Override
//    public void onH264DataCreate(byte[] data, int offset, int length, long presentationTimeUs, int flag) {
//        if(frameListener != null){
//            frameListener.onFrame(data, offset, length, presentationTimeUs, flag);
//        }
//    }

    public void setFrameListener(FrameListener frameListener) {
        this.frameListener = frameListener;
    }

    public int getVideoDen() {
        return videoDen;
    }

    public int getVideoNum() {
        return videoNum;
    }
}
